import struct
import socket
'''
源端口号[0:15]   目的端口号[16:31]
序号
确认序号
首部长度[0:3] 保留[4:9] URG ACK PSH PST SYN FIN 窗口大小[16:31]
检验和[0:15] 紧急指针[16:31]
选项部分
数据
'''
class TCP:
    def __init__(self,src_ip,dst_ip,weikeep,protocol_id,total_len,src_port,dst_port,seq,ack,head_len,keep,URG,ACK,PSH,RST,SYN,FIN,windowSize,checkSum,UrgPtr,other,data):
        self.src_ip = socket.inet_aton(src_ip)
        self.dst_ip = socket.inet_aton(dst_ip)
        self.weikeep= weikeep
        self.protocol_id = protocol_id
        self.total_len = total_len
        self.src_port = src_port  #源端口 2字节 unsignedshort "H"
        self.dst_port = dst_port  #目的端口 2字节 unsignedshort "H"
        self.seq = seq            #序号 4字节 unsignedLong "L"
        self.ack = ack            #确认号 4字节 unsignedint "L"

        self.head_len = head_len  #首部长度 4位
        self.keep = keep          #保留字段，默认全0，6位
        self.URG = URG
        self.ACK = ACK #
        self.PSH = PSH #
        self.RST = RST #
        self.SYN = SYN #
        self.FIN = FIN #          -->两字节 unsignedshort "H"
        self.windowSize = windowSize #窗口大小 两字节 unsignedshort "H"
        self.checkSum = checkSum     #首部检验和 两字节 unsignedshort "H"
        self.UrgPtr = UrgPtr         #紧急指针 两字节 unsignedshort "H"
        self.other = other           #选项部分 可无
        self.data = data             #数据部分
    

    def weiheaderPack(self):
        weiheader = struct.pack("!4s4sBBH",self.src_ip,self.dst_ip,self.weikeep,self.protocol_id,self.total_len)
        return weiheader

    
    def getHeaderPack(self):
        self.flags = self.FIN + (self.SYN << 1) + (self.RST << 2) + (self.PSH <<3) + (self.ACK << 4) + (self.URG << 5)
        keep = self.keep<<6
        head_len = self.head_len<<12
        header = struct.pack("!HHLLHHHH",self.src_port,self.dst_port,self.seq,self.ack,
                    head_len+keep+self.flags,self.windowSize,self.checkSum,
                    self.UrgPtr)
        return header
    

    def getPack(self):
        header = self.getHeaderPack()
        self.total_len = len(header+bytes(self.data,encoding="utf-8"))
        weiheader = self.weiheaderPack()
        res = weiheader+header+bytes(self.data,encoding="utf-8")
        if len(res) % 2 != 0:
	        res += b'\0'
        self.checkSum = self.makeCheckSum(res)
        header = self.getHeaderPack()


        return header+bytes(self.data,encoding="utf-8")
    

    def makeCheckSum(self,msg):
        s = 0
        for i in range(0,len(msg),2):
            w = (msg[i] << 8) + (msg[i+1])
            s = s+w
        s = (s>>16) + (s & 0xffff)
        s = ~s & 0xffff
        return s


if __name__ == "__main__":
    myTCP = TCP(src_ip="192.168.1.5",
                dst_ip="192.168.1.2",
                weikeep=0,
                protocol_id=6,
                total_len=0,
                src_port=1000,
                dst_port=80,
                seq=0,
                ack=0,
                head_len=5,
                keep=0,
                URG=0,
                ACK=0,
                PSH=0,
                RST=0,
                SYN=1,
                FIN=0,
                windowSize=3000,
                checkSum=0,
                UrgPtr=0,
                other=None,
                data="appl"
                )
    pack = myTCP.getPack()
    print(pack)